#include "kernel_operator.h"

using namespace AscendC;

constexpr int32_t BUFFER_SIZE = 256; // UB缓冲区大小

class KernelGelu {
public:
    __aicore__ inline KernelGelu() {}

    // 初始化函数
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, uint32_t totalLength) {
        this->x = x;
        this->y = y;
        this->totalLength = totalLength;
        
        this->blockLength = totalLength / GetBlockNum();
        this->tileNum = this->blockLength / BUFFER_SIZE;
        
        xBuf.SetBuffer(BUFFER_SIZE);
        yBuf.SetBuffer(BUFFER_SIZE);
    }

    // GELU计算核心
    __aicore__ inline void GeluCompute(LocalTensor<float>& output, 
                                     const LocalTensor<float>& input, 
                                     uint32_t len) {
        constexpr float kAlpha = M_2_SQRTPI * M_SQRT1_2;
        constexpr float kBeta = 0.044715f;
        
        // 计算立方项
        LocalTensor<float> cubeBuf;
        cubeBuf.SetBuffer(BUFFER_SIZE);
        Cube(cubeBuf, input, len);

        // 多项式项计算
        LocalTensor<float> termBuf;
        termBuf.SetBuffer(BUFFER_SIZE);
        Muls(termBuf, cubeBuf, kBeta, len);
        Add(termBuf, input, termBuf, len);

        // 缩放项
        LocalTensor<float> scaledBuf;
        scaledBuf.SetBuffer(BUFFER_SIZE);
        Muls(scaledBuf, termBuf, kAlpha, len);

        // tanh激活
        Tanh(scaledBuf, scaledBuf, len);

        // 最终组合
        LocalTensor<float> oneBuf;
        oneBuf.SetBuffer(BUFFER_SIZE);
        Muls(oneBuf, scaledBuf, 0.5f, len);
        Add(oneBuf, oneBuf, 0.5f, len);
        Mul(output, input, oneBuf, len);
    }

    __aicore__ inline void Process() {
        for(uint32_t i = 0; i < this->tileNum; ++i) {
            DataCopy(xBuf, this->x + i * BUFFER_SIZE, BUFFER_SIZE);
            GeluCompute(yBuf, xBuf, BUFFER_SIZE);
            DataCopy(this->y + i * BUFFER_SIZE, yBuf, BUFFER_SIZE);
        }
    }

private:
    GM_ADDR x, y;
    uint32_t totalLength;
    uint32_t blockLength;
    uint32_t tileNum;
    LocalTensor<float> xBuf, yBuf;
};

extern "C" __global__ __aicore__ void gelu_custom(GM_ADDR x, GM_ADDR y, uint32_t totalLength) {
    KernelGelu op;
    op.Init(x, y, totalLength);
    op.Process();
}